import numpy as np
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt


class Qnet(torch.nn.Module):
    ''' 只有一层隐藏层的Q网络 '''
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(Qnet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)
    
class VAnet(torch.nn.Module):
    ''' 只有一层隐藏层的A网络和V网络 '''
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(VAnet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  # 共享网络部分
        self.fc_A = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_V = torch.nn.Linear(hidden_dim, 1)

    def forward(self, x):
        A = self.fc_A(F.relu(self.fc1(x)))
        V = self.fc_V(F.relu(self.fc1(x)))
        # 动作状态价值函数Q = 状态价值函数V + 优势函数A - 优势函数A的平均值
        Q = V + A - A.mean()
        return Q


class DQN:
    ''' DQN算法,包括Double DQN和Dueling DQN '''
    def __init__(self,state_dim,hidden_dim,action_dim,learning_rate,discount_factor,action_epsilon,target_update,device,dqn_type='DuelingDoubleDQN'):
        self.action_dim = action_dim
        
        if dqn_type == 'DuelingDQN' or dqn_type == 'DuelingDoubleDQN':  # Dueling DQN采取不一样的网络框架
            self.q_net = VAnet(state_dim, hidden_dim,self.action_dim).to(device)
            self.target_q_net = VAnet(state_dim, hidden_dim,self.action_dim).to(device)
        else: # VanillaDQN
            self.q_net = Qnet(state_dim, hidden_dim,self.action_dim).to(device)
            self.target_q_net = Qnet(state_dim, hidden_dim,self.action_dim).to(device)
            
        self.optimizer = torch.optim.Adam(self.q_net.parameters(),lr=learning_rate)
        self.discount_factor = discount_factor
        self.action_epsilon = action_epsilon
        self.target_update = target_update
        self.count = 0
        self.dqn_type = dqn_type
        self.device = device

    def take_action_epsilon_greedy(self, state):
        if np.random.random() < self.action_epsilon:
            action = np.random.randint(self.action_dim)
        else:
            state = torch.tensor(state, dtype=torch.float).to(self.device)
            action = self.q_net(state).argmax().item()
        return action
    
    def take_action(self, state):
        state = torch.tensor(state, dtype=torch.float).to(self.device)
        action = self.q_net(state).argmax().item()
        return action

    def max_q_value(self, state):
        state = torch.tensor(state, dtype=torch.float).to(self.device)
        return self.q_net(state).max().item()

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],dtype=torch.float).view(-1, 1).to(self.device)
        actions = actions.long() # int32 → int64

        q_values = self.q_net(states).gather(1, actions)
        # max_next_q_values: next_states 下个状态的最大Q值
        if self.dqn_type == 'DoubleDQN' or self.dqn_type == 'DuelingDoubleDQN': 
            # DoubleDQN 中，选取使得训练Q网络的Q值最大的动作，然后求该动作对应的目标Q网络的Q值
            max_action = self.q_net(next_states).max(1)[1].view(-1, 1)
            max_next_q_values = self.target_q_net(next_states).gather(1, max_action)
        else: 
            # DQN 中，选取使得目标Q网络的Q值最大的动作，然后求该动作对应的目标Q网络的Q值
            max_next_q_values = self.target_q_net(next_states).max(1)[0].view(-1, 1)
        q_targets = rewards + self.discount_factor * max_next_q_values * (1 - dones) # TD误差目标
        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets)) # 均方误差损失函数
        self.optimizer.zero_grad() # PyTorch中默认梯度会累积,这里需要显式将梯度置为0
        dqn_loss.backward()
        self.optimizer.step()

        if self.count % self.target_update == 0: # 更新目标网络
            self.target_q_net.load_state_dict(self.q_net.state_dict())
        self.count += 1

def plot_Qvalue(max_q_value_list, env_name, alg_name, string_train_test = 'Training'):
    plt.plot(list(range(len(max_q_value_list))), max_q_value_list)
    plt.axhline(0, c='orange', ls='--')
    # plt.axhline(10, c='red', ls='--')
    plt.xlabel('Frames')
    plt.ylabel('Q value')
    plt.title(f'{alg_name} on {env_name} ({string_train_test})')
    plt.show()



